#include "lua.h"
#include "lauxlib.h"
#include "lua_scheduler.h"
#include "cm_os.h"
#include "app_log.h"

#define MAX_TASKS 16
#define INSTRUCTIONS_PER_YIELD 5000

static lua_State *g_main_L = NULL;
static lua_task_t g_task_list[MAX_TASKS];
static int g_current_task = -1;

// The hook function that forces a yield
static void yield_hook(lua_State *L, lua_Debug *ar) {
    lua_yield(L, 0);
}

void lua_scheduler_init(lua_State *L) {
    g_main_L = L;
    for (int i = 0; i < MAX_TASKS; i++) {
        g_task_list[i].co = NULL;
        g_task_list[i].status = TASK_DEAD;
        g_task_list[i].ref = LUA_NOREF;
        g_task_list[i].wakeup_time = 0;
    }
}

lua_task_t* lua_scheduler_get_current_task(void) {
    if (g_current_task >= 0 && g_current_task < MAX_TASKS) {
        return &g_task_list[g_current_task];
    }
    return NULL;
}

int lua_scheduler_new_task(lua_State *L) {
    if (!lua_isfunction(L, 1)) {
        return luaL_error(L, "argument to sys.task must be a function");
    }

    int free_slot = -1;
    for (int i = 0; i < MAX_TASKS; i++) {
        if (g_task_list[i].status == TASK_DEAD) {
            free_slot = i;
            break;
        }
    }

    if (free_slot == -1) {
        return luaL_error(L, "max number of tasks reached");
    }

    lua_State *co = lua_newthread(L);

    // Keep a reference to the thread so it's not garbage collected
    g_task_list[free_slot].ref = luaL_ref(L, LUA_REGISTRYINDEX);
    g_task_list[free_slot].co = co;
    g_task_list[free_slot].status = TASK_SUSPENDED;
    g_task_list[free_slot].wakeup_time = 0;

    // Move function and arguments to the new coroutine
    lua_xmove(L, co, lua_gettop(L));

    return 0; // No value is returned to the calling Lua script
}

void lua_scheduler_run(void) {
    // This is a simple round-robin scheduler
    g_current_task = (g_current_task + 1) % MAX_TASKS;

    lua_task_t *current_task = &g_task_list[g_current_task];

    if (current_task->status == TASK_SUSPENDED) {
        // Check if the task is sleeping
        if (current_task->wakeup_time > 0 && osKernelGetTickCount() < current_task->wakeup_time) {
            return; // Still sleeping, skip to next task
        }
        current_task->wakeup_time = 0; // Reset wakeup time

        lua_State *co = current_task->co;

        // Set the hook before resuming
        lua_sethook(co, yield_hook, LUA_MASKCOUNT, INSTRUCTIONS_PER_YIELD);

        current_task->status = TASK_RUNNING;

        int status = lua_resume(co, g_main_L, 0);

        // After resume, the task could be us, or it could be a different one if a hook yielded
        // We always set the returned task to suspended.
        g_task_list[g_current_task].status = TASK_SUSPENDED;

        if (status == LUA_OK) {
            // Coroutine finished its execution
            current_task->status = TASK_DEAD;
            luaL_unref(g_main_L, LUA_REGISTRYINDEX, current_task->ref);
            current_task->ref = LUA_NOREF;
        } else if (status != LUA_YIELD) {
            // An error occurred
            const char* err_msg = lua_tostring(co, -1);
            app_log_printf("[E/LUA_SCHED] Task %d error: %s", g_current_task, err_msg);
            current_task->status = TASK_DEAD;
            luaL_unref(g_main_L, LUA_REGISTRYINDEX, current_task->ref);
            current_task->ref = LUA_NOREF;
        }
    }
}
